package mac5789.projeto4.algorithm;

import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.Map;
import java.util.Set;

import mac5789.projeto4.model.Action;
import mac5789.projeto4.model.State;
import mac5789.projeto4.model.TransitionModel;
import mac5789.projeto4.model.WorldModel;

/**
 * Trial em um RealTimeMDP
 * 
 * @author lundberg
 * 
 */
public class RTDPMDPTrial {

	private WorldModel world;

	private BellmanBackup backup;

	private LinkedList<State> path = new LinkedList<State>();

	private Map<State, Double> stateValue = new HashMap<State, Double>();

	private Set<State> visited = new HashSet<State>();

	/**
	 * Construtor
	 * 
	 * @param world world
	 * @param backup backup
	 */
	public RTDPMDPTrial(WorldModel world, BellmanBackup backup) {
		this.world = world;
		this.backup = backup;
	}

	/**
	 * Executa a melhoria do bellman backup anterior, utilizando o conhecimento
	 * adquirido no trial.
	 * 
	 * @return bellman backup
	 */
	public BellmanBackup improve() {
		return this.backup.improve(this.path);
	}

	/**
	 * Executa o trial. Deve ser chamado uma única vez.
	 */
	public void execute() {
		assert path.isEmpty();
		path.add(world.initialState());
		visited.add(world.initialState());
		while (!world.isGoal(currentState()) && !isMovingInCircles()) {
			this.move();
		}
	}

	private boolean isMovingInCircles() {
		return this.path.size() > this.visited.size();
	}

	/**
	 * Efetua um movimento no trial, escolhendo a próxima ação e definindo o
	 * próximo estado.
	 */
	private void move() {
		Action action = this.pickAction();
		update(action);

		TransitionModel transitionModel = world.transitionModelFor(currentState(), action);
		if (visited.containsAll(transitionModel.possibleStates())) {
			State nextState = transitionModel.pickState();
			path.add(nextState);
			visited.add(nextState);
		} else {
			State nextState;
			do {
				nextState = transitionModel.pickState();
			} while (visited.contains(nextState));
			path.add(nextState);
			visited.add(nextState);
		}
	}

	/**
	 * Atualiza o valor do último estado, dada a ação que foi escolhida
	 * 
	 * @param action ação escolhida
	 */
	private void update(Action action) {
		this.stateValue.put(currentState(), this.stateQValue(currentState(), action));
	}

	/**
	 * Escolhe a próxima ação que deve ser executada no trial
	 * 
	 * @return ação escolhida
	 */
	private Action pickAction() {
		double min = Double.POSITIVE_INFINITY;
		Action action = null;
		for (Action candidate : world.availableActionsAt(currentState())) {
			double stateQValue = this.stateQValue(currentState(), candidate);
			if (stateQValue < min) {
				min = stateQValue;
				action = candidate;
			}
		}
		return action;
	}

	/**
	 * Calcula do QValue do estado e da ação, dado pelo custo da ação mais os
	 * valores dos estados provenientes de sua execução, com pesos proporcionais
	 * às probabilidades
	 * 
	 * @param state state
	 * @param action action
	 * @return QValue
	 */
	private double stateQValue(State state, Action action) {
		double qValue = world.rewardModel().cost(action);
		TransitionModel transitionModel = world.transitionModelFor(state, action);
		for (State next : transitionModel.possibleStates()) {
			qValue += transitionModel.probability(next) * this.stateValue(next);
		}
		return qValue;
	}

	/**
	 * Valor do estado. Caso não tenha sido calculado pelo próprio Trial, busca
	 * o valor do backup
	 * 
	 * @param state state
	 * @return valor do estado
	 */
	private double stateValue(State state) {
		if (!stateValue.containsKey(state)) {
			stateValue.put(state, -backup.valueFor(state));
		}
		return stateValue.get(state);
	}

	/**
	 * Busca o estado atual do trial
	 * 
	 * @return estado
	 */
	private State currentState() {
		return path.getLast();
	}
}
